﻿#ifndef __RBT__
#define __RBT__
#include <iostream>
/**
* 红黑树
* 作者：邵明博
* 实现了对红黑树的基本操作：插入、删除、搜索、遍历以及选择
*/
template <class item>
class RBT{
public :
	RBT();
	~RBT();
	void insert(item key);
	void remove(item key);
	bool search(item key);
	void traverse();
	item select(int n);
	bool empty();
private:
	struct node{
		item data;
		node *left, *right;
		int child_count;//该字段需要在插入，旋转以及删除(联合左右子树)操作中维护
		bool red;
		node(item data):data(data),left(NULL),right(NULL),child_count(0),red(true){}
	};
	typedef node * link;
	link root;//字段 红黑树树根

	/*内部方法实现*/
	void destroy(link&);//析构的内部实现
	bool search_in(item, link); //搜索方法的内部实现
	void remove_in(item, link&); // 删除方法的内部实现
	void insert_in(item, link&, int); //插入方法的内部实现
	void traverse_in(std::ostream&, link);//遍历输出树中的每个结点
	item select_in(int, link);//选取第n 小的数据项
	void right_rotate(link&);//向右旋转
	void left_rotate(link&);//向左旋转
	void part(int, link&);//将第n 小的数据项作为树根
	link join_left_and_right(link& left, link& right){
		if(right == NULL) return left;
		part(1, right);
		right->left = left;
		right->child_count += left ? left->child_count + 1 : 0;
		return right;
	}//在删除操作中，联合被删除结点的左右两颗子树
	/*内部方法实现*/
};
template <class item>
void RBT<item>::destroy(link& root){
	if(root == NULL) return;
	destroy(root->left); 
	destroy(root->right);
	delete root; root = NULL;
}

template <class item>
bool RBT<item>::search_in(item key, link root){
	if(root == NULL) return false;
	if(key == root->data) return true;
	else if(key < root->data) return search_in(key, root->left);
	else return search_in(key, root->right);
}

template <class item>
void RBT<item>::traverse_in(std::ostream& out, link root){
	if(root == NULL) return ;
	out << root->data << " ";
	traverse_in(out, root->left);
	traverse_in(out, root->right);
}

template <class item>
void RBT<item>::right_rotate(link& node){
	//旋转后child_count 字段的变化
	int add = node->right ? node->right->child_count + 1 + 1: 1 ;
	int del = node->left ? (node->left->left ? node->left->left->child_count + 1 + 1 : 1) : 1;
	node->child_count -= del; 	node->left->child_count += add;
	//旋转
	link x = node->left; node->left = x->right; x->right = node; node = x;
}

template <class item>
void RBT<item>::left_rotate(link& node){
	//旋转后child_count 字段的变化
	int add = node->left ? node->left->child_count + 1 + 1 : 1 ;
	int del = node->right ? (node->right->right ? node->right->right->child_count + 1 + 1 : 1) : 1;
	node->child_count -= del; 	node->right->child_count += add;
	//旋转
	link x = node->right; node->right = x->left; x->left = node; node = x;
}

template <class item>
item RBT<item>::select_in(int n, link root){
	assert(n >1 && root != NULL);
	int left_child_count = root->left ? root->left->child_count + 1 : 0;
	if(left_child_count+1 == n) return root->data;
	else if(n < left_child_count+1) return select_in(n, root->left);
	else return select_in(n - left_child_count - 1, root->right);
}

template <class item>
void RBT<item>::part(int n, link& node){
	if(node == NULL || n < 1) return;
	int left_child_count = node->left ? node->left->child_count + 1 : 0;
	if(n == left_child_count + 1) return ;
	else if(n < left_child_count + 1) { part(n, node->left); right_rotate(node); }
	else { part(n - left_child_count - 1, node->right); left_rotate(node); }
}

/**
* 方法：红黑树的删除内部实现
* 功能描述：
*	1.搜索待删除的项目，若未找到，直接返回；
*   2.若找到了，则将左子树和右子树合并: 在右子树中选取待删除项的后继，将左子树链接上去。
*   3.回溯过程中，对搜索路径中的每一项的child_count - 1
*/
template <class item>
void RBT<item>::remove_in(item key, link& root){
	if(root == NULL) return ;
	static bool found = false;
	if(key < root->data) { remove_in(key, root->left); if(found) --root->child_count; }
	else if(key > root->data) { remove_in(key, root->right); if(found) --root->child_count; }
	else{
		found = true;
		link del = root;
		root = join_left_and_right(root->left, root->right);
		delete del;
	}
}

/*
* 方法:	红黑树的插入
* 功能描述：
*1.每次插入一个红色结点作为叶子结点，但根结点作为黑色结点
*2.每次插入向下检查：若包含2个儿子结点，并且均为红色结点，则将儿子结点变为黑色，父节点变为红色
*3.插入后，向上回溯，检查以下2种情况：
*  a.两个指向红色结点的链接在同一个方向（eg:root->left->red == root->left->left->red == true）,
*	 进行一次旋转（right_rotate(root)），后改变新root 的颜色为black, root->right->red 为 true
*  b.两个指向红色结点的链接不在同一个方向（eg:root->red == root->left->red == true）,进行一次旋转
*	 （right_rotate(root)）. 
*  a,b 两种情况出现在不同的抽象维度，所以才用layer 来进行区分：通常是向右子树insert 之后（交叉回溯），
*  才会出现b 的情况。而a 的情况：有可能是b 旋转之后得到的结果，也有可能是直接插入后的直接回溯结果。 
*/
template <class item>
void RBT<item>::insert_in(item key, link& root, int layer){
	static bool not_found = true;
	if(root == NULL) { root = new node(key); return; }
	if(root->left && root->left->red && root->right && root->right->red){
		root->red = true; root->left->red = false; root->right->red = false;
	}
	if(key < root->data){
		insert_in(key, root->left, 0);
		if(not_found) ++root->child_count;
		if(root && root->red && root->left && root->left->red && layer) right_rotate(root);
		if(root->left && root->left->red && root->left->left && root->left->left->red){
			right_rotate(root); root->red = false; root->right->red = true;
		}
	}else if(key > root->data){
		insert_in(key, root->right, 1); 
		if(not_found) ++root->child_count;
		if(root && root->red && root->right && root->right->red && !layer) left_rotate(root);
		if(root->right && root->right->red && root->right->right && root->right->right->red){
			left_rotate(root); root->red = false; root->left->red = true;
		}
	}else not_found = false;
}


/*********************************************
               外部接口
********************************************/
template <class item>
RBT<item>::RBT(){ root = NULL; }

template <class item>
RBT<item>::~RBT(){ 	destroy(root); }

template <class item>
bool RBT<item>::empty(){ return root == NULL; }

template <class item>
bool RBT<item>::search(item key){ return search_in(key, root); }

template <class item>
void RBT<item>::traverse(){ traverse_in(std::cout, root); }

template <class item>
item RBT<item>::select(int n) { return select_in(n, root); }

template <class item>
void RBT<item>::insert(item key) { insert_in(key, root, 0); root->red = false; }

template <class item>
void RBT<item>::remove(item key) { remove_in(key, root); }
#endif